from tri_training import *
from learn import *
import numpy as np
import sklearn
import scipy.io as scio
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from Graphic import *

def loadData():
    dataFile = './norb.mat'
    data = scio.loadmat(dataFile)

    traindata = np.double(data['train_x'] / 255)
    trainlabel = np.double(data['train_y'])
    testdata = np.double(data['test_x'] / 255)
    testlabel = np.double(data['test_y'])

    data = np.row_stack([traindata, testdata])  # 行合并操作
    label = np.row_stack([trainlabel, testlabel]).argmax(axis=1)

    train_index = np.random.choice(data.shape[0], 8600, replace=False) # 选8600个
    rest_index = list(set(np.arange(data.shape[0])) - set(train_index))
    test_index = np.random.choice(rest_index, 10000, replace=False) # 选10000个
    u_index = list(set(rest_index) - set(test_index))

    traindata = data[train_index]
    trainlabel = label[train_index]

    testdata = data[test_index]
    testlabel = label[test_index]

    udata = data[u_index]

    print(traindata.shape, testdata.shape, udata.shape)
    print(trainlabel.shape, testlabel.shape)

    return traindata, trainlabel, udata, testdata, testlabel


def loadDataNews():
    fileNameTrain = '20newsGroup_train.txt'
    fileNameTrainLabel = '20newsGroup_train_label.txt'
    fileNameTest = '20newsGroup_test.txt'
    fileNameTestLabell = '20newsGroup_test_label.txt'

    traindata = pd.read_csv(fileNameTrain, encoding='utf-8', sep='\t', header=0)  # 如果数据间是标准化格式用“\t”,否则用"\s+"
    trainlabel = pd.read_csv(fileNameTrainLabel, encoding='utf-8', sep='\t', header=0)
    testdata = pd.read_csv(fileNameTest, encoding='utf-8', sep='\t', header=0)
    testlabel = pd.read_csv(fileNameTestLabell, encoding='utf-8', sep='\t', header=0)

    data = np.row_stack([traindata, testdata])  # 行合并操作
    labelTemp = np.row_stack([trainlabel, testlabel])

    label = []
    for num in range(np.shape(labelTemp)[0]):
        label.extend(labelTemp[num])
    label = np.array(label)

    train_index = np.random.choice(data.shape[0], 842, replace=False)  # 选8600个
    rest_index = list(set(np.arange(data.shape[0])) - set(train_index))
    test_index = np.random.choice(rest_index, 999, replace=False)  # 选10000个
    u_index = list(set(rest_index) - set(test_index))

    traindata = data[train_index]
    trainlabel = label[train_index]

    testdata = data[test_index]
    testlabel = label[test_index]

    udata = data[u_index]

    print(traindata.shape, testdata.shape, udata.shape)
    print(trainlabel.shape, testlabel.shape)

    return traindata, trainlabel, udata, testdata, testlabel


if __name__ == '__main__':
    # traindata, trainlabel, udata, testdata, testlabel = loadData()

    # 20news数据集
    traindata, trainlabel, udata, testdata, testlabel = loadDataNews()

    clf = RandomForestClassifier()
    clf.fit(traindata, trainlabel)
    res1 = clf.predict(testdata)
    print(accuracy_score(res1, testlabel))

    TT = TriTraining([RandomForestClassifier(), RandomForestClassifier(), RandomForestClassifier()])
    TT.fit(traindata, trainlabel, udata)
    res2 = TT.predict(testdata)
    print(accuracy_score(res2, testlabel))


    # changeMat = TT.getChange()
    # print(changeMat)
    # G1 = Graphic(changeMat[0])
    # G1.display()
    # G2 = Graphic(changeMat[1])
    # G2.display()
    # G3 = Graphic(changeMat[2])
    # G3.display()

    # learn = Learn()

    # svm = learn.SVM()
    # svm.fit(traindata, trainlabel)
    # resSvm = svm.predict(testdata)
    # print(accuracy_score(resSvm, testlabel))

    # NB = learn.Naive_Bayes()
    # NB.fit(traindata, trainlabel)
    # resNB = NB.predict(testdata)
    # print(accuracy_score(resNB, testlabel))
    #
    # tree = learn.Tree()
    # tree.fit(traindata, trainlabel)
    # resTree = tree.predict(testdata)
    # print(accuracy_score(resTree, testlabel))
    #
    # KNN = learn.KNN()
    # KNN.fit(traindata, trainlabel)
    # resKNN = KNN.predict(testdata)
    # print(accuracy_score(resKNN, testlabel))
    # #
    # TT = TriTraining([KNN, NB, tree])
    # TT.fit(traindata, trainlabel, udata)
    # res2 = TT.predict(testdata)
    # print(accuracy_score(res2, testlabel))